#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File         :pos_control_offb.py
@Description  :位置环控制，主要用于领机的offb控制
@Time         :2024/05/30 16:54:00
@Author       :Lin Yuheng
@Version      :1.0
'''

import rospy
import tty, termios
import sys, select
from mavros_msgs.msg import AttitudeTarget, PositionTarget, State
from mavros_msgs.srv import CommandBool, SetMode
from geometry_msgs.msg import Pose, PoseStamped, Twist, TwistStamped
from nav_msgs.msg import Odometry
from std_msgs.msg import String
import time
import PyKDL
import math
import numpy as np


class Controller:
    def __init__(self):
        self.plane_id = 0
        self.vehicle_id = "plane_0"
        self.px = 0
        self.py = 0
        self.pz = 0
        self.vx, self.vy, self.vz = 0, 0, 0
        self.roll, self.pitch, self.yaw = 0, 0, 0
        self.phi_dot, self.theta_dot, self.psi_dot = 0, 0, 0
        self.send_flag = False

        # subscribers
        self.local_pose_sub = rospy.Subscriber(
            self.vehicle_id + "/mavros/local_position/pose",
            PoseStamped,
            callback=self.local_pose_callback,
        )
        self.mavros_sub = rospy.Subscriber(
            self.vehicle_id + "/mavros/state",
            State,
            callback=self.mavros_state_callback,
        )

        # pub
        self.target_motion_pub = rospy.Publisher(
            self.vehicle_id + "/mavros/setpoint_position/local",
            PoseStamped,
            queue_size=10
        )

        self.target_raw_pose = PoseStamped()

    def local_pose_callback(self, msg: PoseStamped):
        self.px = msg.pose.position.x
        self.py = msg.pose.position.y
        self.pz = msg.pose.position.z

    def local_vel_callback(self, msg: TwistStamped):
        self.vx = msg.twist.linear.x
        self.vy = msg.twist.linear.y
        self.vz = msg.twist.linear.z

    def odom_callback(self, msg: Odometry):
        x = msg.pose.pose.orientation.x
        y = msg.pose.pose.orientation.y
        z = msg.pose.pose.orientation.z
        w = msg.pose.pose.orientation.w
        rot = PyKDL.Rotation.Quaternion(x, y, z, w)
        self.roll = rot.GetRPY()[0]
        self.pitch = rot.GetRPY()[1]
        self.yaw = rot.GetRPY()[2]

        self.phi_dot = msg.twist.twist.angular.x
        self.theta_dot = msg.twist.twist.angular.y
        self.psi_dot = msg.twist.twist.angular.z

    def mavros_state_callback(self, msg: State):
        self.mavros_state = msg.mode
        self.arm_state = msg.armed

    def start(self):
        print("start fixed-wing position control")
        rospy.init_node(
            self.vehicle_id + "_position_control_node"
        )
        rate = rospy.Rate(50)
        wp_idx = 0
        # wps_list = [[0, 120], [200, 120], [200, -120], [0, -120]] # ENU坐标系
        wps_list = [[0, 150], [200, 150], [200, -100], [0, -100]] # ENU坐标系

        while not rospy.is_shutdown():
            dist = math.sqrt(
                (self.px - wps_list[wp_idx][0]) ** 2
                + (self.py - wps_list[wp_idx][1]) ** 2
            )
            if dist < 30:
                wp_idx += 1
                wp_idx = wp_idx % 4

            east = wps_list[wp_idx][0]
            north = wps_list[wp_idx][1]
            upward = 50

            self.target_raw_pose.pose.position.x = east
            self.target_raw_pose.pose.position.y = north
            self.target_raw_pose.pose.position.z = upward

            try:
                self.target_motion_pub.publish(self.target_raw_pose)
            except rospy.ROSException as e:
                pass
            rate.sleep()

            rospy.loginfo(
                "wp:%d; position: (%f, %f); distance:%f", wp_idx+1, self.px, self.py, dist
            )


if __name__ == "__main__":
    Con = Controller()
    try:
        Con.start()
    except rospy.ROSInterruptException:
        pass
